# K = 3 and d = 4, m = 8, n = 6, C^(A) = 4, C^(B) = 2
import numpy as np
import time

m = 16
n = 8
CA = 4
CB = 4

# k = 1, d = 1
A_1_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_1_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 1, d = 2
A_1_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_1_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 1, d = 3
A_1_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_1_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 1, d = 4
A_1_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_1_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 2, d = 1
A_2_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_2_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 2, d = 2
A_2_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_2_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 2, d = 3
A_2_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_2_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 2, d = 4
A_2_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_2_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 3, d = 1
A_3_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_3_1 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 3, d = 2
A_3_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_3_2 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 3, d = 3
A_3_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_3_3 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# k = 3, d = 4
A_3_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (m, CA))
B_3_4 = np.random.normal(0, 1/np.sqrt(2*np.pi), (n, CB))

# x
X = np.random.uniform(0.0, 1.0, (m,m,m,m))

# W
W = np.random.normal(0, 1/np.sqrt(2*np.pi), (m*m*m*m, n*n*n*n))


def normal_algo_1():
	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_1_3 = np.kron(A_1_3, B_1_3)
	AB_1_4 = np.kron(A_1_4, B_1_4)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_2_3 = np.kron(A_2_3, B_2_3)
	AB_2_4 = np.kron(A_2_4, B_2_4)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)
	AB_3_3 = np.kron(A_3_3, B_3_3)
	AB_3_4 = np.kron(A_3_4, B_3_4)

	W_1 = np.concatenate([AB_1_1, AB_2_1, AB_3_1], axis = -1)
	W_2 = np.concatenate([AB_1_2, AB_2_2, AB_3_2], axis = -1)
	W_3 = np.concatenate([AB_1_3, AB_2_3, AB_3_3], axis = -1)
	W_4 = np.concatenate([AB_1_4, AB_2_4, AB_3_4], axis = -1)

	# W_1 to (m_1 * n_1*C)
	W_1 = np.reshape(np.reshape(W_1, (m, n, -1)), (m, -1))

	# Y is (m_2*m_3*m_4 * n_1*C)
	Y = np.matmul(np.reshape(X, (m, m*m*m)).T, W_1)

	# Y to (n_1*m_3*m_4 * m_2*C)
	Y = np.reshape(Y, (m, m*m, n, -1))
	Y = np.transpose(Y, (2, 1, 0, 3))
	Y = np.reshape(Y, (n*m*m, -1))

	# W_2 to (m_2*C * n_2)
	W_2 = np.reshape(W_2, (m, n, -1))
	W_2 = np.transpose(W_2, (0, 2, 1))
	W_2 = np.reshape(W_2, (-1, n))

	# Y is (n_1*m_3*m_4 * n_2)
	Y = np.matmul(Y, W_2)

	# Y to (n_1*n_2*m_4 * m_3)
	Y = np.reshape(Y, (n, m, m, n))
	Y = np.transpose(Y, (0, 3, 2, 1))
	Y = np.reshape(Y, (n*n*m, m))

	# W_3 to (m_3 * n_3*C)
	W_3 = np.reshape(np.reshape(W_3, (m, n, -1)), (m, -1))

	# Y is (n_1*n_2*m_4 * n_3*C)
	Y = np.matmul(Y, W_3)

	# Y to (n_1*n_2*n_3 * m_4*C)
	Y = np.reshape(Y, (n*n, m, n, -1))
	Y = np.transpose(Y, (0, 2, 1, 3))
	Y = np.reshape(Y, (n*n*n, -1))

	# W_4 to (m_4*C * n_4)
	W_4 = np.reshape(W_4, (m, n, -1))
	W_4 = np.transpose(W_4, (0, 2, 1))
	W_4 = np.reshape(W_4, (-1, n))

	# Y is (n_1*n_2*n_3 * n_4)
	Y = np.matmul(Y, W_4)

	return np.reshape(Y, (-1))


def divided_algo_1():
	def single_algo_1(I, W_1, W_2):
		# W_1 to (m_1 * n_1*C)
		W_1 = np.reshape(np.reshape(W_1, (m, n, -1)), (m, -1))

		# O is (m_2*m_3*m_4 * n_1*C)
		O = np.matmul(np.reshape(I, (m, m*m*m)).T, W_1)

		# O to (n_1*m_3*m_4 * m_2*C)
		O = np.reshape(O, (m, m*m, n, -1))
		O = np.transpose(O, (2, 1, 0, 3))
		O = np.reshape(O, (n*m*m, -1))

		# W_2 to (m_2*C * n_2)
		W_2 = np.reshape(W_2, (m, n, -1))
		W_2 = np.transpose(W_2, (0, 2, 1))
		W_2 = np.reshape(W_2, (-1, n))

		# O is (n_1*m_3*m_4 * n_2)
		O = np.matmul(O, W_2)

		# O to (n_1*n_2*m_4, m_3)
		O = np.reshape(O, (n, m, m, n))
		O = np.transpose(O, (0, 3, 2, 1))
		O = np.reshape(O, (n*n*m, m))

		return O

	def single_algo_1_2(O, W_3, W_4):
		# W_3 to (m_3 * n_3*C)
		W_3 = np.reshape(np.reshape(W_3, (m, n, -1)), (m, -1))

		# O is (n_1*n_2*m_4, n_3*C)
		O = np.matmul(O, W_3)

		# O to (n_1*n_2*n_3, m_4*C)
		O = np.reshape(O, (n*n, m, n, -1))
		O = np.transpose(O, (0, 2, 1, 3))
		O = np.reshape(O, (n*n*n, -1))

		# W_4 to (m_4*C * n_4)
		W_4 = np.reshape(W_4, (m, n, -1))
		W_4 = np.transpose(W_4, (0, 2, 1))
		W_4 = np.reshape(W_4, (-1, n))

		# O is (n_1*n_2*n_3, n_4)
		O = np.matmul(O, W_4)

		return np.reshape(O, (-1))

	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_1_3 = np.kron(A_1_3, B_1_3)
	AB_1_4 = np.kron(A_1_4, B_1_4)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_2_3 = np.kron(A_2_3, B_2_3)
	AB_2_4 = np.kron(A_2_4, B_2_4)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)
	AB_3_3 = np.kron(A_3_3, B_3_3)
	AB_3_4 = np.kron(A_3_4, B_3_4)

	X_1 = single_algo_1(X, AB_1_1, AB_1_2)
	X_2 = single_algo_1(X, AB_2_1, AB_2_2)
	X_3 = single_algo_1(X, AB_3_1, AB_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_1_2(Y, AB_1_3, AB_1_4)
	X_2 = single_algo_1_2(Y, AB_2_3, AB_2_4)
	X_3 = single_algo_1_2(Y, AB_3_3, AB_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_2():
	def single_algo_2(I, A_1, B_1, A_2, B_2):
		# O is (m_2*m_3*m_4 * C_A)
		O = np.matmul(np.reshape(I, (m, m*m*m)).T, A_1)

		# O is (m_2*m_3*m_4*n_1 * C_A*C_B)
		O = np.kron(O, B_1)

		# O is (n_1*m_3*m_4*C_B * m_2*C_A)
		O = np.reshape(O, (m, -1, n, CA, CB))
		O = np.transpose(O, (2, 1, 4, 0, 3))
		O = np.reshape(O, (-1, m*CA))

		# O is (n_1*m_3*m_4 * C_B)
		O = np.matmul(O, np.reshape(A_2, (-1)))
		O = np.reshape(O, (-1, CB))

		# O is (n_1*n_2*m_3*m_4)
		O = np.matmul(O, B_2.T)
		O = np.reshape(O, (n, -1, n))
		O = np.reshape(np.transpose(O, (0, 2, 1)), (-1))

		return O

	def single_algo_2_2(O, A_3, B_3, A_4, B_4):
		# O to (n_1*n_2*m_4 * m_3)
		O = np.reshape(O, (-1, m, m))
		O = np.reshape(np.transpose(O, (0, 2, 1)), (-1, m))

		# O is (n_1*n_2*m_4 * C_A)
		O = np.matmul(O, A_3)

		# O is (n_1*n_2*m_4*n_3 * C_A*C_B)
		O = np.kron(O, B_3)

		# O is (n_1*n_2*n_3*C_B * m_4*C_A)
		O = np.reshape(O, (-1, m, n, CA, CB))
		O = np.transpose(O, (0, 2, 4, 1, 3))
		O = np.reshape(O, (-1, m*CA))

		# O is (n_1*n_2*n_3 * C_B)
		O = np.matmul(O, np.reshape(A_4, (-1)))
		O = np.reshape(O, (-1, CB))

		# O is (n_1*n_2*n_3*n_4)
		O = np.matmul(O, B_4.T)
		O = np.reshape(O, (-1))
		
		return O

	X_1 = single_algo_2(X, A_1_1, B_1_1, A_1_2, B_1_2)
	X_2 = single_algo_2(X, A_2_1, B_2_1, A_2_2, B_2_2)
	X_3 = single_algo_2(X, A_3_1, B_3_1, A_3_2, B_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_2_2(Y, A_1_3, B_1_3, A_1_4, B_1_4)
	X_2 = single_algo_2_2(Y, A_2_3, B_2_3, A_2_4, B_2_4)
	X_3 = single_algo_2_2(Y, A_3_3, B_3_3, A_3_4, B_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_1_reverse(XR):
	def single_algo_1(I, W_1, W_2):
		# W_1 to (m_1 * n_1*C)
		W_1 = np.reshape(np.reshape(W_1, (m, n, -1)), (m, -1))

		# O is (m_4*m_3*m_2 * n_1*C)
		O = np.matmul(np.reshape(I, (m*m*m, m)), W_1)

		# O is (n_1*m_4*m_3 * m_2*C)
		O = np.reshape(O, (-1, m, n, CA*CB))
		O = np.transpose(O, (2, 0, 1, 3))
		O = np.reshape(O, (-1, m*CA*CB))

		# W_2 to (m_2*C * n_2)
		W_2 = np.reshape(W_2, (m, n, -1))
		W_2 = np.transpose(W_2, (0, 2, 1))
		W_2 = np.reshape(W_2, (-1, n))

		# O is (n_2*n_1*m_4 * m_3)
		O = np.reshape(np.matmul(O, W_2).T, (-1, m))

		return O

	def single_algo_1_2(O, W_3, W_4):
		# W_3 to (m_3 * n_3*C)
		W_3 = np.reshape(np.reshape(W_3, (m, n, -1)), (m, -1))

		# O is (n_2*n_1*m_4 * n_3*C)
		O = np.matmul(O, W_3)

		# O is (n_3*n_2*n_1 * m_4*C)
		O = np.reshape(O, (-1, m, n, CA*CB))
		O = np.transpose(O, (2, 0, 1, 3))
		O = np.reshape(O, (-1, m*CA*CB))

		# W_4 to (m_4*C * n_4)
		W_4 = np.reshape(W_4, (m, n, -1))
		W_4 = np.transpose(W_4, (0, 2, 1))
		W_4 = np.reshape(W_4, (-1, n))

		# O is (n_4*n_3*n_2*n_1)
		O = np.reshape(np.matmul(O, W_4).T, (-1))

		return O

	AB_1_1 = np.kron(A_1_1, B_1_1)
	AB_1_2 = np.kron(A_1_2, B_1_2)
	AB_1_3 = np.kron(A_1_3, B_1_3)
	AB_1_4 = np.kron(A_1_4, B_1_4)
	AB_2_1 = np.kron(A_2_1, B_2_1)
	AB_2_2 = np.kron(A_2_2, B_2_2)
	AB_2_3 = np.kron(A_2_3, B_2_3)
	AB_2_4 = np.kron(A_2_4, B_2_4)
	AB_3_1 = np.kron(A_3_1, B_3_1)
	AB_3_2 = np.kron(A_3_2, B_3_2)
	AB_3_3 = np.kron(A_3_3, B_3_3)
	AB_3_4 = np.kron(A_3_4, B_3_4)

	X_1 = single_algo_1(XR, AB_1_1, AB_1_2)
	X_2 = single_algo_1(XR, AB_2_1, AB_2_2)
	X_3 = single_algo_1(XR, AB_3_1, AB_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_1_2(Y, AB_1_3, AB_1_4)
	X_2 = single_algo_1_2(Y, AB_2_3, AB_2_4)
	X_3 = single_algo_1_2(Y, AB_3_3, AB_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_2_reverse(XR):
	def single_algo_2(I, A_1, B_1, A_2, B_2):
		# O is (m_4*m_3*m_2 * C_A)
		O = np.matmul(np.reshape(I, (m*m*m, m)), A_1)

		# O is (n_1*m_4*m_3*m_2 * C_B*C_A)
		O = np.kron(B_1, O)

		# O is (n_1*m_4*m_3*C_B * m_2*C_A)
		O = np.reshape(O, (-1, m, CB, CA))
		O = np.transpose(O, (0, 2, 1, 3))
		O = np.reshape(O, (-1, m*CA))

		# O is (n_1*m_4*m_3 * C_B)
		O = np.matmul(O, np.reshape(A_2, (-1)))
		O = np.reshape(O, (-1, CB))

		# O is (n_2*n_1*m_4 * m_3)
		O = np.matmul(O, B_2.T).T
		O = np.reshape(np.reshape(O, (-1)), (-1, m))

		return O

	def single_algo_2_2(O, A_3, B_3, A_4, B_4):
		# O is (n_2*n_1*m_4 * C_A)
		O = np.matmul(O, A_3)

		# O is (n_3*n_2*n_1*m_4 * C_B*C_A)
		O = np.kron(B_3, O)

		# O is (n_3*n_2*n_1*C_B * m_4*C_A)
		O = np.reshape(O, (-1, m, CB, CA))
		O = np.transpose(O, (0, 2, 1, 3))
		O = np.reshape(O, (-1, m*CA))

		# O is (n_3*n_2*n_1 * C_B)
		O = np.matmul(O, np.reshape(A_4, (-1)))
		O = np.reshape(O, (-1, CB))

		# O is (n_4*n_3*n_2*n_1)
		O = np.reshape(np.matmul(O, B_4.T).T, (-1))

		return O

	X_1 = single_algo_2(XR, A_1_1, B_1_1, A_1_2, B_1_2)
	X_2 = single_algo_2(XR, A_2_1, B_2_1, A_2_2, B_2_2)
	X_3 = single_algo_2(XR, A_3_1, B_3_1, A_3_2, B_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_2_2(Y, A_1_3, B_1_3, A_1_4, B_1_4)
	X_2 = single_algo_2_2(Y, A_2_3, B_2_3, A_2_4, B_2_4)
	X_3 = single_algo_2_2(Y, A_3_3, B_3_3, A_3_4, B_3_4)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_new1_notransp():
	# A is (m_1*m_2) B is (n_1*n_2)
	def single_algo_1(I, A, B):
		# O is (m_3*m_4)
		O = np.matmul(np.reshape(A, (-1)), np.reshape(I, (m*m, -1)))

		# O is (m_3*m_4 * n_1*n_2)
		O = np.outer(O, np.reshape(B, (-1)))

		return O

	# A is (m_3*m_4) B is (n_3*n_4)
	def single_algo_1_2(O, A, B):
		# O is (n_1*n_2)
		O = np.matmul(np.reshape(A, (-1)), O)

		# O is (n_1*n_2 * n_3*n_4)
		O = np.outer(O, np.reshape(B, (-1)))

		return np.reshape(O, (-1))

	# k = 1
	A_1_12 = np.matmul(A_1_1, A_1_2.T)
	B_1_12 = np.matmul(B_1_1, B_1_2.T)
	A_1_34 = np.matmul(A_1_3, A_1_4.T)
	B_1_34 = np.matmul(B_1_3, B_1_4.T)

	# k = 2
	A_2_12 = np.matmul(A_2_1, A_2_2.T)
	B_2_12 = np.matmul(B_2_1, B_2_2.T)
	A_2_34 = np.matmul(A_2_3, A_2_4.T)
	B_2_34 = np.matmul(B_2_3, B_2_4.T)

	# k = 3
	A_3_12 = np.matmul(A_3_1, A_3_2.T)
	B_3_12 = np.matmul(B_3_1, B_3_2.T)
	A_3_34 = np.matmul(A_3_3, A_3_4.T)
	B_3_34 = np.matmul(B_3_3, B_3_4.T)

	X_1 = single_algo_1(X, A_1_12, B_1_12)
	X_2 = single_algo_1(X, A_2_12, B_2_12)
	X_3 = single_algo_1(X, A_3_12, B_3_12)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_1_2(Y, A_1_34, B_1_34)
	X_2 = single_algo_1_2(Y, A_2_34, B_2_34)
	X_3 = single_algo_1_2(Y, A_3_34, B_3_34)
	Y = X_1 + X_2 + X_3

	return Y


def divided_algo_new2_notransp():
	def single_algo_2(I, A_1, B_1, A_2, B_2):
		# O is (C_A * m_2*m_3*m_4)
		O = np.matmul(A_1.T, np.reshape(I, (m, -1)))

		# O is (C_A*m_2*m_3*m_4*n_1*C_B)
		O = np.outer(np.reshape(O, (-1)), np.reshape(B_1, (-1)))
		O = np.reshape(O, (-1))

		# O is (m_3*m_4*n_1*C_B)
		O = np.matmul(np.reshape(A_2.T, (-1)), np.reshape(O, (CA*m, -1)))

		# O is (m_3*m_4*n_1*n_2)
		O = np.matmul(np.reshape(O, (-1, CB)), B_2.T)
		O = np.reshape(O, (-1))

		return O

	def single_algo_2_2(I, A_3, B_3, A_4, B_4):
		# O is (C_A * m_4*n_1*n_2)
		O = np.matmul(A_3.T, np.reshape(I, (m, -1)))

		# O is (C_A*m_4*n_1*n_2*n_3*C_B)
		O = np.outer(np.reshape(O, (-1)), np.reshape(B_3, (-1)))
		O = np.reshape(O, (-1))

		# O is (n_1*n_2*n_3*C_B)
		O = np.matmul(np.reshape(A_4.T, (-1)), np.reshape(O, (CA*m, -1)))

		# O is (n_1*n_2*n_3*n_4)
		O = np.matmul(np.reshape(O, (-1, CB)), B_4.T)
		O = np.reshape(O, (-1))

		return O

	X_1 = single_algo_2(X, A_1_1, B_1_1, A_1_2, B_1_2)
	X_2 = single_algo_2(X, A_2_1, B_2_1, A_2_2, B_2_2)
	X_3 = single_algo_2(X, A_3_1, B_3_1, A_3_2, B_3_2)
	Y = X_1 + X_2 + X_3

	X_1 = single_algo_2_2(Y, A_1_3, B_1_3, A_1_4, B_1_4)
	X_2 = single_algo_2_2(Y, A_2_3, B_2_3, A_2_4, B_2_4)
	X_3 = single_algo_2_2(Y, A_3_3, B_3_3, A_3_4, B_3_4)
	Y = X_1 + X_2 + X_3

	return Y


if __name__ == '__main__':
	s0 = time.time()
	Y_0 = np.matmul(np.reshape(X, (-1)),W)
	e0 = time.time()

	s1 = time.time()
	Y_1 = normal_algo_1()
	e1 = time.time()

	s2 = time.time()
	Y_2 = divided_algo_1()
	e2 = time.time()

	s3 = time.time()
	Y_3 = divided_algo_2()
	e3 = time.time()

	XR = np.reshape(np.transpose(np.reshape(X, (m, m, m, m)), (3, 2, 1, 0)), (-1))

	s4 = time.time()
	Y_4 = divided_algo_1_reverse(XR)
	e4 = time.time()
	Y_4 = np.reshape(np.transpose(np.reshape(Y_4, (n, n, n, n)), (3, 2, 1, 0)), (-1))

	s5 = time.time()
	Y_5 = divided_algo_2_reverse(XR)
	e5 = time.time()
	Y_5 = np.reshape(np.transpose(np.reshape(Y_5, (n, n, n, n)), (3, 2, 1, 0)), (-1))

	s6 = time.time()
	Y_6 = divided_algo_new1_notransp()
	e6 = time.time()

	s7 = time.time()
	Y_7 = divided_algo_new2_notransp()
	e7 = time.time()

	print('Vector Matrix (%lf) is:' % ((e0 - s0) * 1000))
	print(Y_0.shape)
	print(Y_0)
	print('\n')

	print('Normal (%lf) is:' % ((e1 - s1) * 1000))
	print(Y_1.shape)
	print(Y_1)
	print('\n')

	print('Divided 1 (%lf) is:' % ((e2 - s2) * 1000))
	print(Y_2.shape)
	print(Y_2)
	print('\n')

	print('Divided 2 (%lf) is:' % ((e3 - s3) * 1000))
	print(Y_3.shape)
	print(Y_3)
	print('\n')

	print('Divided 1 reversing (%lf) is:' % ((e4 - s4) * 1000))
	print(Y_4.shape)
	print(Y_4)
	print('\n')

	print('Divided 2 reversing (%lf) is:' % ((e5 - s5) * 1000))
	print(Y_5.shape)
	print(Y_5)
	print('\n')

	print('Divided 1 no transp (%lf) is:' % ((e6 - s6) * 1000))
	print(Y_6.shape)
	print(Y_6)
	print('\n')

	print('Divided 2 no transp (%lf) is:' % ((e7 - s7) * 1000))
	print(Y_7.shape)
	print(Y_7)
	print('\n')
